#!/usr/bin/env python

"""

A Phased Array Processor Interface for the SMA
written by Rurik Primiani
with help from...

Aaron Parsons: who wrote ibob_serial, which has been modified here
   to include a close() member function to free up the serial port.
   The style of the commented class titles have also been taken from
   Aaron's python codes :)
   
"""

from __future__ import division
from math import floor, pi
from time import sleep, time
from socket import socket, AF_INET, SOCK_STREAM
from pylab import plot, subplot, show, clf, xlim, ylim
from numpy import max
from ibob_serial import IBOBSerial
from paUtils import *

# A pylab color macro
colors = ['b','g', 'r', 'c', 'm', 'y', 'k', 'w']

#  ____                    ____
# |  _ \  _____  ___    __/ ___}
# | | \ \/ __/ |/   \  / / /
# | | | / __/| | /_\ \/ / /
# | |/ / /__ | |_ _  | |\ \___
# |___/\____}|___} |_|_| \____}
#

class DelayCorrector(IBOBSerial):
    def __init__(self, port):
        self.comport = port
        self.interface = {  #'CH0'    : (int('D0002000', 16), 'b'),
                            #'CH1'    : (int('D0004000', 16), 'b'),
                            #'CH2'    : (int('D0006000', 16), 'b'),
                            #'CH3'    : (int('D0008000', 16), 'b'),
                            #'SUM'    : (int('D3002000', 16), 'b'),
                            #'bias0'  : (int('D0000000', 16), 'L'),
                            #'bias1'  : (int('D0000100', 16), 'L'),
                            #'bias2'  : (int('D0000200', 16), 'L'),
                            #'bias3'  : (int('D0000300', 16), 'L'),
                            #'c'      : (int('D0000800', 16), 'L'),
                            'cnt_to' : (int('D0000400', 16), 'L'),
                            'count'  : (int('D0000500', 16), 'L'),
                            'delay0' : (int('D0000600', 16), 'L'),
                            'delay1' : (int('D0000700', 16), 'L'),
                            'delay2' : (int('D0000800', 16), 'L'),
                            'delay3' : (int('D0000900', 16), 'L'),
                            'fdelay0': (int('D0000A00', 16), 'L'),
                            'fdelay1': (int('D0000B00', 16), 'L'),
                            'fdelay2': (int('D0000C00', 16), 'L'),
                            'fdelay3': (int('D1000000', 16), 'L'),
                            'h0_00'  : (int('D1000100', 16), 'L'),
                            'h0_01'  : (int('D1000200', 16), 'L'),
                            'h0_10'  : (int('D1000300', 16), 'L'),
                            'h0_11'  : (int('D1000400', 16), 'L'),
                            'h1_00'  : (int('D1000500', 16), 'L'),
                            'h1_01'  : (int('D1000600', 16), 'L'),
                            'h1_10'  : (int('D1000700', 16), 'L'),
                            'h1_11'  : (int('D1000800', 16), 'L'),
                            'h2_00'  : (int('D1000900', 16), 'L'),
                            'h2_01'  : (int('D1000A00', 16), 'L'),
                            'h2_10'  : (int('D1000B00', 16), 'L'),
                            'h2_11'  : (int('D1000C00', 16), 'L'),
                            'h3_00'  : (int('D1000D00', 16), 'L'),
                            'h3_01'  : (int('D1000E00', 16), 'L'),
                            'h3_10'  : (int('D1000F00', 16), 'L'),
                            'h3_11'  : (int('D1001000', 16), 'L'),
                            'his_en' : (int('D2000000', 16), 'L'),
                            'his_rst': (int('D2000100', 16), 'L'),
                            'reset'  : (int('D2000200', 16), 'L'),
                            #'reset_c': (int('D2000C00', 16), 'L'),
                            'select0': (int('D2000300', 16), 'L'),
                            'select1': (int('D2000400', 16), 'L'),
                            'select2': (int('D2000500', 16), 'L'),
                            'select3': (int('D2000600', 16), 'L'),
                            'source' : (int('D2000700', 16), 'L'),
                            'xaui'   : (int('D2000800', 16), 'L'),
                            'thresh0': (int('D2000900', 16), 'L'),
                            'thresh1': (int('D2000A00', 16), 'L'),
                            'thresh2': (int('D2000B00', 16), 'L'),
                            'thresh3': (int('D2000C00', 16), 'L'),
                            'valid'  : (int('D2000D00', 16), 'L'),
                            'xaui0f' : (int('D2000E00', 16), 'L'),
                            'xaui1f' : (int('D2000F00', 16), 'L'), }
        self.params = {'delay0'  : 4,
                       'delay1'  : 4,
                       'delay2'  : 4,
                       'delay3'  : 4,
                       'fdelay0' : 0,
                       'fdelay1' : 0,
                       'fdelay2' : 0,
                       'fdelay3' : 0,
                       'select0' : 0,
                       'select1' : 0,
                       'select2' : 0,
                       'select3' : 0,
                       #'bias0'   : 0,
                       #'bias1'   : 0,
                       #'bias2'   : 0,
                       #'bias3'   : 0,
                       'thresh0' : 16,
                       'thresh1' : 16,
                       'thresh2' : 16,
                       'thresh3' : 16, }
        self.cntrls = {#'c'       : 0,
                       #'reset_c' : 0,
                       'reset'   : 0,
                       'his_en'  : 0,
                       'his_rst' : 0,
                       'source'  : 1,
                       'xaui'    : 1,
                       'valid'   : 15,
                       'cnt_to'  : 1024000000-1, }
        self.hist0 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h0_cnt = 0
        self.hist1 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h1_cnt = 0
        self.hist2 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h2_cnt = 0
        self.hist3 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h3_cnt = 0
        self.ch0 = ()
        self.ch1 = ()
        self.ch2 = ()
        self.ch3 = ()
        self.sum = ()
        self.devices = {}
        self.types = {}
        for i in self.interface:
            self.devices[i], self.types[i] = self.interface[i]
        self.ibob = IBOBSerial(self.devices, self.types, port)

    def read(self, dev, num):
        val = self.ibob.read(dev,num)
        return val

    def write(self, dev, num, size=1):
        self.ibob.write(dev,num)
        if self.ibob.read(dev,size)==num:
            return None
        else:
            return "Err: Readback failed"

    def start(self):
        for p in self.params:
            self.ibob.write(p,self.params[p])
        for c in self.cntrls:
            self.ibob.write(c,self.cntrls[c])

    #def record(self):
    #    self.ibob.write('reset_c',1)
    #    self.ibob.write('reset_c',0)
    #    self.ibob.write('c',1)
    #    sleep(0.01)
    #    self.ibob.write('c',0)
    #    self.ch0 = self.ibob.read('CH0',8000)
    #    self.ch1 = self.ibob.read('CH1',8000)
    #    self.ch2 = self.ibob.read('CH2',8000)
    #    self.ch3 = self.ibob.read('CH3',8000)

    def getHistogram(self,intime):
        self.ibob.write('his_rst',1)
        self.ibob.write('his_rst',0)
        self.ibob.write('his_en',1)
        sleep(intime)
        self.ibob.write('his_en',0)
        self.h0_cnt = 0
        self.h1_cnt = 0
        self.h2_cnt = 0
        self.h3_cnt = 0
        for i in self.hist0:
            self.hist0[i] = self.ibob.read('h0_'+i)
            self.h0_cnt += self.hist0[i]
        for i in self.hist1:
            self.hist1[i] = self.ibob.read('h1_'+i)
            self.h1_cnt += self.hist1[i]
        for i in self.hist2:
            self.hist2[i] = self.ibob.read('h2_'+i)
            self.h2_cnt += self.hist2[i]
        for i in self.hist3:
            self.hist3[i] = self.ibob.read('h3_'+i)
            self.h3_cnt += self.hist3[i]

    def setDelays(self, delay={0:0,1:0,2:0,3:0}):
        for d in delay:
            de = 'delay'+str(d)
            fe = 'fdelay'+str(d)
            se = 'select'+str(d)
            self.params[de]  = floor(floor(delay[d]/10)/4)+4
            self.params[fe] = floor((delay[d]-(floor(delay[d]/10)*10)))
            self.params[se] = 3-floor(delay[d]/10)%4
            self.ibob.write(de,self.params[de])
            self.ibob.write(fe,self.params[fe])
            self.ibob.write(se,self.params[se])

    def setThresholds(self, threshold={0:16,1:16,2:16,3:16}):
        for t in threshold:
            th = 'thresh'+str(t)
            self.params[th] = threshold[t]
            self.ibob.write(th,self.params[th])

    def setBiases(self, bias={0:0,1:0,2:0,3:0}):
        for b in bias:
            bi = 'bias'+str(b)
            self.params[bi] = bias[b]
            self.ibob.write(bi,self.params[bi])

    def close(self):
        self.ibob.close_ibob()

#
#
#
# BEE2 Cross Correlator
#
#
#

class CrossCorrelator:
    def __init__(self, host="sma-bee2.cfa.harvard.edu", port=4500):
        self.host = host
        self.port = port

    def cmd(self, command):
        size = 1024

        s = socket(AF_INET, SOCK_STREAM)
        s.connect((self.host,self.port))
        s.send(command)
        result = s.recv(size)
        s.close()

        return result
    
    def read_reg(self, reg):
        val = self.cmd('read_reg:'+reg)
        return int(val)

    def read_corr(self, bram):
        strval = self.cmd('read_bram:'+bram+':4:32')
        val = []
        for i in strval.split(','):
                val.append(int(i))
        return val

    def correlate(self, int_time):
        self.cmd('correlate:'+str(int_time))

    def getCorrFunc(self):
        baselines = []
        for j in range(7):
            bram = "corr_out"+str(j)
            val = self.read_corr(bram)
            baselines.append(val)
        return baselines
    
#  ____
# |  _  \___
# | |_| / _ \
# |  __/ /_\ \
# |_|  |_| |_|
#

class PhasedArray(delayCorrector, crossCorrelator):
    def __init__(self, dec1="COM4", dec2="COM3", dbe="COM5", bee2_host="sma-bee2.cfa.harvard.edu", bee2_port=4500):
        self.dec1 = delayCorrector(dec1)
        self.dec2 = delayCorrector(dec2)
        self.corr = crossCorrelator(bee2_host,bee2_port)
        self.hist0 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h0_cnt = 0
        self.hist1 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h1_cnt = 0
        self.hist2 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h2_cnt = 0
        self.hist3 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h3_cnt = 0
        self.hist4 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h4_cnt = 0
        self.hist5 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h5_cnt = 0
        self.hist6 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h6_cnt = 0
        self.hist7 = { '01'      : 0,
                       '00'      : 0,
                       '10'      : 0,
                       '11'      : 0, }
        self.h7_cnt = 0
        self.correlations = []
        
    def start(self):
        self.dec1.start()
        self.dec2.start()

    def getHistograms(self):
        self.dec1.ibob.write('his_rst',1)
        self.dec2.ibob.write('his_rst',1)
        self.dec1.ibob.write('his_rst',0)
        self.dec2.ibob.write('his_rst',0)
        self.dec1.ibob.write('his_en',1)
        self.dec2.ibob.write('his_en',1)
        sleep(5)
        self.dec1.ibob.write('his_en',0)
        self.dec2.ibob.write('his_en',0)
        self.h0_cnt = 0
        self.h1_cnt = 0
        self.h2_cnt = 0
        self.h3_cnt = 0
        self.h4_cnt = 0
        self.h5_cnt = 0
        self.h6_cnt = 0
        self.h7_cnt = 0
        for i in self.hist0:
            self.dec1.hist0[i] = self.dec1.ibob.read('h0_'+i)
            self.hist0[i] = self.dec1.ibob.read('h0_'+i)
            self.h0_cnt += self.hist0[i]
        for i in self.hist1:
            self.dec1.hist1[i] = self.dec1.ibob.read('h1_'+i)
            self.hist1[i] = self.dec1.ibob.read('h1_'+i)
            self.h1_cnt += self.hist1[i]
        for i in self.hist2:
            self.dec1.hist2[i] = self.dec1.ibob.read('h2_'+i)
            self.hist2[i] = self.dec1.ibob.read('h2_'+i)
            self.h2_cnt += self.hist2[i]
        for i in self.hist3:
            self.dec1.hist3[i] = self.dec1.ibob.read('h3_'+i)
            self.hist3[i] = self.dec1.ibob.read('h3_'+i)
            self.h3_cnt += self.hist3[i]
        for i in self.hist4:
            self.dec2.hist0[i] = self.dec2.ibob.read('h0_'+i)
            self.hist4[i] = self.dec2.ibob.read('h0_'+i)
            self.h4_cnt += self.hist4[i]
        for i in self.hist5:
            self.dec2.hist1[i] = self.dec2.ibob.read('h1_'+i)
            self.hist5[i] = self.dec2.ibob.read('h1_'+i)
            self.h5_cnt += self.hist5[i]
        for i in self.hist6:
            self.dec2.hist2[i] = self.dec2.ibob.read('h2_'+i)
            self.hist6[i] = self.dec1.ibob.read('h2_'+i)
            self.h6_cnt += self.hist6[i]
        for i in self.hist7:
            self.dec2.hist3[i] = self.dec2.ibob.read('h3_'+i)
            self.hist7[i] = self.dec2.ibob.read('h3_'+i)
            self.h7_cnt += self.hist7[i]

    def getCorrelations(self, int_time):
        self.corr.correlate(int_time)
        sleep(int_time)
        self.correlations = self.corr.getCorrFunc()

    def plotCorrelations(self):
        clf()
        for i in range(len(self.correlations)):
            corr = self.correlations[i]
            subplot(len(self.correlations),1,i+1)
            plot(range(-16,16), corr, colors[i]+"-", linewidth=2)
            xlim([-16,15])
        show()
            
    def plotPhases(self):
        clf()
        for i in range(len(self.correlations)):
            phase = getPhase(corrReorder(self.correlations[i]))
            subplot(len(self.correlations),1,i+1)
            plot(range(len(phase)), phase, colors[i]+"h", markersize=5)
            ylim([-pi,pi])
        show()

    def plotMagnitudes(self):
        clf()
        for i in range(len(self.correlations)):
            mag = getMagnitude(corrReorder(self.correlations[i]))
            subplot(len(self.correlations),1,i+1)
            plot(range(len(mag)), mag, colors[i]+"-", linewidth=2)
        show()

    def plotBaselines(self):
        clf()
        for i in range(len(self.correlations)):
            mag = getMagnitude(corrReorder(self.correlations[i]))
            phase = getPhase(corrReorder(self.correlations[i]))
            subplot(len(self.correlations),1,i+1)
            plot(range(len(phase)), phase, colors[i]+"h", markersize=5)
            plot(range(len(mag)), pi*array(mag/max(mag)), colors[i]+"-", linewidth=2)
            ylim([-pi*1.25,pi*1.25])
        show()

    def stop(self):
        self.dec1.close()
        self.dec2.close()
